/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
    treeHeight = 0;
    inited = false;
    writable = false;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
    PageId endPid = -1;
    RC r = RC_ERROR;

    // Open the index file.
    if ((r = pf.open(indexname, mode)) != 0)
        return r;

    // Check if the index file is newly created.
    endPid = pf.endPid();
    if (endPid == 0) {
        // The index file is newly created. So initialize file header, and
        // create the first node ever, which will be leaf node AND root node.
        BTLeafNode root;
        BTPageByteBuffer pbb;
        rootPid = 1;
        treeHeight = 1;
        if (!pbb.appendInt(rootPid)) {
            pf.close();
            return RC_ERROR;
        }
        if (!pbb.appendInt(treeHeight)) {
            pf.close();
            return RC_ERROR;
        }
        if ((r = pf.write(0, pbb.rawBuffer())) != 0) {
            pf.close();
            return r;
        }
        if (pf.endPid() != 1) {
            pf.close();
            return RC_ERROR;
        }
        if ((r = root.write(1, pf)) != 0) {
            pf.close();
            return r;
        }
    }
    else if (endPid > 0){
        // Load the file header.
        char buffer[PageFile::PAGE_SIZE];
        if ((r = pf.read(0, &buffer)) != 0) {
            pf.close();
            return r;
        }
        BTPageByteBuffer pbb(buffer);
        if (!pbb.readInt(rootPid)) {
            pf.close();
            return RC_ERROR;
        }
        if (!pbb.readInt(treeHeight)) {
            pf.close();
            return RC_ERROR;
        }
        if (rootPid <= 0 || treeHeight <= 0) {
            pf.close();
            return RC_INVALID_FILE_FORMAT;
        }
    }
    else {
        pf.close();
        return RC_ERROR;
    }

    // Initialization complete!
    inited = true;
    if (mode == 'w' || mode == 'W')
        writable = true;
    else
        writable = false;
    return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    RC r = RC_ERROR;
    if ((r = pf.close()) != 0)
        return r;
    inited = false;
    writable = false;
    return 0;
}

/*
 * Internal helper function for BTreeIndex::_insert_into_leaf().
 * @param pid1[IN] a pointer to all the nodes less than the key
 *                 (used only when need to create new root).
 * @param key[IN] the key to be inserted into the parent.
 * @param pid2[IN] a pointer to all the nodes greater than the key.
 * @param parentPids[IN/OUT] a stack containing the parents of the current
 *                           node, up to the root.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::_insert_into_parent(PageId pid1, int key, PageId pid2, stack<PageId> &parentPids)
{
    if (parentPids.empty()) {
        BTNonLeafNode newRoot;
        RC r = RC_ERROR;

        if ((r = newRoot.initializeRoot(pid1, key, pid2)) != 0)
            return r;

        if ((r = newRoot.write(rootPid = pf.endPid(), pf)) != 0)
            return r;

        treeHeight++;

        BTPageByteBuffer pbb;
        if (!pbb.appendInt(rootPid))
            return RC_ERROR;
        if (!pbb.appendInt(treeHeight))
            return RC_ERROR;
        if ((r = pf.write(0, pbb.rawBuffer())) != 0)
            return r;

        return 0;
    }
    else {
        BTNonLeafNode parent;
        PageId parentPid = -1;
        RC r = RC_ERROR;

        parentPid = parentPids.top();
        parentPids.pop();

        if ((r = parent.read(parentPid, pf)) != 0)
            return r;

        if ((r = parent.insert(key, pid2)) != 0) {
            if (r == RC_NODE_FULL) {
                BTNonLeafNode sibling;
                int midKey = 0;
                PageId siblingPid = -1;
                if ((r = parent.insertAndSplit(key, pid2, sibling, midKey)) != 0)
                    return r;
                if ((r = parent.write(parentPid, pf)) != 0)
                    return r;
                if ((r = sibling.write(siblingPid = pf.endPid(), pf)) != 0)
                    return r;
                return _insert_into_parent(parentPid, midKey, siblingPid, parentPids);
            }
            else
                return r;
        }
        else
            return parent.write(parentPid, pf);
    }
}

/*
 * Internal helper function for BTreeIndex::insert().
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @param level[IN] the current level of the tree (0 ... treeHeight - 1) for
 *                  keeping track of recursion.
 * @param pid[IN] the pid of the node at which to start the recursion
 *                (indicates the current node).
 * @param leafIndex[IN] once at the leaf level, this keeps track of which leaf
 *                      node we are at (leftmost leaf node has index 0).
 * @param parentPids[IN/OUT] a stack containing the parents of the current
 *                           node, up to the root.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::_insert_into_leaf(int key, const RecordId& rid, int level, PageId pid, int leafIndex, stack<PageId> &parentPids)
{
    if (level == treeHeight - 1) {
        // We're at the leaf level.

        BTLeafNode node;
        int eid = -1;
        RC r = RC_ERROR;

        // Load the leaf node.
        if ((r = node.read(pid, pf)) != 0)
            return r;

        // Can we insert into this leaf node?
        if ((r = node.locate(key, eid)) != 0) {
            if (r == RC_NO_SUCH_RECORD) {
                // Try the next leaf node if possible. If a next leaf node does
                // not exist, or the smallest key in the next leaf node is
                // larger than the key to be inserted, then it is okay to
                // insert into the current leaf node.
                PageId nextNode = node.getNextNodePtr();
                if (nextNode >= 0) {
                    r = _insert_into_leaf(key, rid, level, node.getNextNodePtr(), leafIndex + 1, parentPids);
                    if (r == 0 || r != RC_BACKTRACK)
                        return r;
                    // else r == RC_BACKTRACK, so proceed to insert into the
                    // current leaf node.
                }
            }
        }
        else if (leafIndex > 0 && eid == 0){
            int existingKey = 0;
            RecordId existingRid = {-1, -1};
            if ((r = node.readEntry(eid, existingKey, existingRid)) != 0)
                return r;
            if (existingKey > key)
                return RC_BACKTRACK; // Insert into earlier leaf node instead.
            // else proceed to insert into the current leaf node.
        }
        // else proceed to insert into the current leaf node.

        // Insert into the current leaf node.
        if ((r = node.insert(key, rid)) != 0) {
            if (r == RC_NODE_FULL) {
                // The leaf node is full, so split, etc.
                BTLeafNode sibling;
                int siblingKey = 0;
                PageId siblingPid = -1;
                if ((r = node.insertAndSplit(key, rid, sibling, siblingKey)) != 0)
                    return r;
                if ((r = sibling.setNextNodePtr(node.getNextNodePtr())) != 0)
                    return r;
                if ((r = sibling.write(siblingPid = pf.endPid(), pf)) != 0)
                    return r;
                if ((r = node.setNextNodePtr(siblingPid)) != 0)
                    return r;
                if ((r = node.write(pid, pf)) != 0)
                    return r;
                return _insert_into_parent(pid, siblingKey, siblingPid, parentPids);
            }
            else
                return r;
        }
        else
            return node.write(pid, pf);
    }
    else if (level < treeHeight - 1) {
        // We're at a nonleaf level.

        BTNonLeafNode node;
        PageId childPid = -1;
        RC r = RC_ERROR;

        // Load the nonleaf node.
        if ((r = node.read(pid, pf)) != 0)
            return r;

        // Find a pointer in the nonleaf node to follow.
        if ((r = node.locateChildPtr(key, childPid)) != 0)
            return r;

        // Follow the pointer.
        parentPids.push(pid);
        return _insert_into_leaf(key, rid, level + 1, childPid, 0, parentPids);
    }

    return RC_ERROR;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
    if (!inited)
        return RC_UNINITIALIZED;
    if (!writable)
        return RC_READONLY;
    if (rid.pid < 0 || rid.sid < 0)
        return RC_INVALID_RID;
    stack<PageId> parentPids;
    return _insert_into_leaf(key, rid, 0, rootPid, 0, parentPids);
}

/*
 * Internal helper function for BTreeIndex::locate().
 * @param searchKey[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @param level[IN] the current level of the tree (0 ... treeHeight - 1) for
 *                  keeping track of recursion.
 * @param pid[IN] the pid of the node at which to start the recursion.
 *                (indicates the current node).
 * @return error code. 0 if no error.
 */
RC BTreeIndex::_locate(int searchKey, IndexCursor& cursor, int level, PageId pid) const
{
    if (level == treeHeight - 1) {
        // We're at the leaf level.

        BTLeafNode node;
        RC r = RC_ERROR;

        // Load the leaf node.
        if ((r = node.read(pid, pf)) != 0)
            return r;

        // Find an entry in the leaf node.
        if ((r = node.locate(searchKey, cursor.eid)) != 0) {
            if (r == RC_NO_SUCH_RECORD) {
                // An entry was not found in the leaf node. Try searching the
                // sibling leaf node if it exists.
                PageId nextNode = node.getNextNodePtr();
                if (nextNode >= 0)
                    return _locate(searchKey, cursor, level, node.getNextNodePtr());
                else
                    return r;
            }
            else
                return r;
        }

        // An entry has been found in the leaf node. We are done!
        cursor.pid = pid;

        return 0;
    }
    else if (level < treeHeight - 1) {
        // We're at a nonleaf level.

        BTNonLeafNode node;
        PageId childPid = -1;
        RC r = RC_ERROR;

        // Load the nonleaf node.
        if ((r = node.read(pid, pf)) != 0)
            return r;

        // Find a pointer in the nonleaf node to follow.
        if ((r = node.locateChildPtr(searchKey, childPid)) != 0)
            return r;

        // Follow the pointer.
        return _locate(searchKey, cursor, level + 1, childPid);
    }

    return RC_ERROR;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor) const
{
    if (!inited)
        return RC_UNINITIALIZED;
    return _locate(searchKey, cursor, 0, rootPid);
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
    BTLeafNode node;
    RC r = RC_ERROR;

    if (!inited)
        return RC_UNINITIALIZED;

    if (cursor.eid < 0 || cursor.pid < 0)
        return RC_INVALID_CURSOR;

    if ((r = node.read(cursor.pid, pf)) != 0)
        return r;

    if ((r = node.readEntry(cursor.eid, key, rid)) != 0)
        return r;

    cursor.eid++;

    if (cursor.eid >= node.getKeyCount()) {
        cursor.pid = node.getNextNodePtr();
        cursor.eid = 0;
    }

    return 0;
}
